
import cv2
from module.detector import Detector
from utils.tool import *
from onnxsim import simplify
import json

import torch
import numpy as np

model_path = "/project/train/models/mm.pt"
YAML_PATH = "/project/train/FastDet/configs/person.yaml"


class MyModel():
    def __init__(self) -> None:

        self.thresh = 0.65
        self.device = torch.device("cuda")
        self.cfg = LoadYaml(YAML_PATH)
        print("load weight from:%s" % model_path)
        self.model = Detector(self.cfg.category_num, True).to(self.device)
        self.model.load_state_dict(torch.load(
            model_path, map_location=self.device))
        self.model.eval()

    def __call__(self, ori_img):
        # 数据预处理
        res_img = cv2.resize(
            ori_img, (self.cfg.input_width, self.cfg.input_height), interpolation=cv2.INTER_LINEAR)
        img = res_img.reshape(1, self.cfg.input_height,
                              self.cfg.input_width, 3)
        img = torch.from_numpy(img.transpose(0, 3, 1, 2))
        img = img.to(self.device).float() / 255.0
        # 模型推理
        preds = self.model(img)
        # 特征图后处理
        output = handle_preds(preds, self.device, self.thresh)
        # 加载label names
        LABEL_NAMES = ["person"]
        H, W, _ = ori_img.shape
        objs = []
        # 绘制预测框
        for box in output[0]:
            # print(box)
            box = box.tolist()
            obj_score = box[4]
            category = LABEL_NAMES[int(box[5])]
            x1, y1 = int(box[0] * W), int(box[1] * H)
            x2, y2 = int(box[2] * W), int(box[3] * H)
            # cv2.rectangle(ori_img, (x1, y1), (x2, y2), (255, 255, 0), 2)
            # cv2.putText(ori_img, '%.2f' % obj_score,
            #             (x1, y1 - 5), 0, 0.7, (0, 255, 0), 2)
            # cv2.putText(ori_img, category, (x1, y1 - 25),
            #             0, 0.7, (0, 255, 0), 2)
            objs.append([x1, y1, x2-x1, y2-y1, obj_score, category])
        # cv2.imwrite("result.png", ori_img)

        return objs


def init():
    model = MyModel()
    return model


def process_image(handle=None, input_image=None, args=None, ** kwargs):

    objs = handle(input_image)
    print(objs)
    obj_dict = []
    target_info = []
    for x, y, w, h, cf, nm in objs:
        _obj = {
            "x": int(x),
            "y": int(y),
            "width": int(w),
            "height": int(h),
            "confidence": cf,
            "name": nm
        }
        obj_dict.append(_obj)
    fake_result = {}
    fake_result["algorithm_data"] = {
        "is_alert": len(target_info) > 0,
        "target_count": len(target_info),
        "target_info": target_info
    }
    fake_result["model_data"] = {"objects": obj_dict}
    return json.dumps(fake_result, indent=4)


if __name__ == "__main__":
    model_path = "/home/u20/FastestDet/weights/zhatu.pth"
    mode = init()

    # img = cv2.imread("/home/data/599/1014a54.jpg")
    img = cv2.imread(
        "/home/u20/FastestDet/data/person.jpg")
    ans = process_image(mode, img)
    print(ans)
